# utils.py
import os
import sys
import torch
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader, random_split
from torch.optim.lr_scheduler import ReduceLROnPlateau
from tqdm import tqdm

current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if current_dir not in sys.path:
    sys.path.append(current_dir)

from student.time_series import LSTMPredictor
from student.output import DeepResNet
from data.distill_dataset import DistillDataManager



def setup_environment():
    # 获取项目根目录的绝对路径
    current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    if current_dir not in sys.path:
        sys.path.append(current_dir)
    
    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        print(f"当前GPU内存使用: {torch.cuda.memory_allocated(0)/1024**2:.1f} MB")
    
    return device

def load_data(radius_matrix, mode='serial'):
    data_manager = DistillDataManager(radius_matrix)
    
    # 加载训练数据集并分割出验证集
    full_train_dataset = data_manager.get_train_dataset(mode=mode)
    train_size = int(0.8 * len(full_train_dataset))
    val_size = len(full_train_dataset) - train_size
    train_dataset, val_dataset = random_split(full_train_dataset, [train_size, val_size])
    
    # 加载测试数据集
    test_dataset = data_manager.get_test_dataset(mode=mode)
    
    print("训练集大小:", len(train_dataset))
    print("验证集大小:", len(val_dataset))
    print("测试集大小:", len(test_dataset))
    
    return data_manager, train_dataset, val_dataset, test_dataset

def save_metrics(results, filename='test_metrics.npy'):
    np.save(filename, results)
    
    # 打印结果
    print("\n整体性能指标:")
    print(f"均方误差 (MSE): {results['overall']['mse']:.6f}")
    print(f"均方根误差 (RMSE): {results['overall']['rmse']:.6f}")
    print(f"平均绝对误差 (MAE): {results['overall']['mae']:.6f}")
    print(f"R² 分数: {results['overall']['r2']:.6f}")
    
    print("\n各端口性能指标:")
    for port_result in results['ports']:
        port = port_result['port']
        print(f"\n端口 {port}:")
        print(f"  MSE: {port_result['mse']:.6f}")
        print(f"  RMSE: {port_result['rmse']:.6f}")
        print(f"  MAE: {port_result['mae']:.6f}")
        print(f"  R²: {port_result['r2']:.6f}")